from typing import Iterable, List

import torch
import torch.nn as nn


def get_activation(name: str) -> nn.Module:
    name = (name or "relu").lower()
    if name == "relu":
        return nn.ReLU()
    if name == "leakyrelu":
        return nn.LeakyReLU(negative_slope=0.01)
    if name == "gelu":
        return nn.GELU()
    raise ValueError(f"Unsupported activation: {name}")


class MLP(nn.Module):
    def __init__(
        self,
        input_dim: int,
        hidden_dims: Iterable[int],
        num_classes: int,
        activation: str = "relu",
        dropout: float = 0.0,
    ) -> None:
        super().__init__()

        if input_dim <= 0:
            raise ValueError("input_dim must be > 0")
        if num_classes <= 0:
            raise ValueError("num_classes must be > 0")

        hidden_dims_list: List[int] = list(hidden_dims or [])
        act = get_activation(activation)
        layers: List[nn.Module] = []

        prev_dim = input_dim
        for hidden_dim in hidden_dims_list:
            if hidden_dim <= 0:
                raise ValueError("All hidden_dims must be > 0")
            layers.append(nn.Linear(prev_dim, hidden_dim))
            layers.append(act.__class__())  # new instance to avoid shared state
            if dropout and dropout > 0:
                layers.append(nn.Dropout(p=dropout))
            prev_dim = hidden_dim

        layers.append(nn.Linear(prev_dim, num_classes))

        self.network = nn.Sequential(*layers)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.network(x)


